# Copyright (c) Princeton University.
# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.

# Authors: Lahav Lipson, Alexander Raistrick


import bpy
import mathutils
from numpy.random import uniform, normal, randint
from nodes.node_wrangler import Nodes, NodeWrangler
from nodes import node_utils
from nodes.color import color_category
from surfaces import surface
import numpy as np

pallete1 = [   ((0.2632, 0.1493, 0.0558, 1.0), 0.0),
    ((0.2695, 0.1585, 0.064, 1.0), 0.058),
    ((0.829, 0.7488, 0.449, 1.0), 0.0718),
    ((0.3486, 0.223, 0.1024, 1.0), 0.0773),
    ((0.3552, 0.2386, 0.1138, 1.0), 0.2099),
    ((0.5925, 0.4391, 0.2677, 1.0), 0.2486),
    ((0.2278, 0.1492, 0.0705, 1.0), 0.2624),
    ((0.4087, 0.2429, 0.1068, 1.0), 0.3149),
    ((0.3037, 0.1832, 0.0586, 1.0), 0.3204),
    ((0.3279, 0.1901, 0.0636, 1.0), 0.3923),
    ((0.604, 0.4648, 0.2181, 1.0), 0.4006),
    ((0.303, 0.204, 0.0733, 1.0), 0.4227),
    ((0.2477, 0.1717, 0.0419, 1.0), 0.4862),
    ((0.2103, 0.1538, 0.0341, 1.0), 0.547),
    ((0.3617, 0.2418, 0.0928, 1.0), 0.5773),
    ((0.2413, 0.1354, 0.038, 1.0), 0.6215),
    ((0.6602, 0.4602, 0.2262, 1.0), 0.6271),
    ((0.3841, 0.251, 0.1021, 1.0), 0.6685),
    ((0.4158, 0.2773, 0.1234, 1.0), 0.6934),
    ((0.4158, 0.2875, 0.1564, 1.0), 0.7624),
    ((0.7427, 0.5479, 0.3642, 1.0), 0.779),
    ((0.2511, 0.1583, 0.0831, 1.0), 0.8094),
    ((0.6233, 0.4074, 0.2138, 1.0), 0.8149),
    ((0.3241, 0.182, 0.076, 1.0), 0.8453),
    ((0.73, 0.4825, 0.2818, 1.0), 0.8619),
    ((0.7319, 0.5246, 0.3017, 1.0), 0.884),
    ((0.2873, 0.1812, 0.0778, 1.0), 0.9088),
    ((0.5537, 0.381, 0.2168, 1.0), 0.9309),
    ((0.3119, 0.1916, 0.0883, 1.0), 0.9558),
    ((0.931, 0.7854, 0.5187, 1.0), 0.9613),
    ((0.9281, 0.7416, 0.4916, 1.0), 1.0)]

pallete2 = [   ((0.3586, 0.3253, 0.1285, 1.0), 0.0),
    ((0.2603, 0.2268, 0.0725, 1.0), 0.0319),
    ((0.2972, 0.2467, 0.0864, 1.0), 0.0957),
    ((0.4891, 0.4373, 0.1874, 1.0), 0.117),
    ((0.3481, 0.3054, 0.1156, 1.0), 0.1383),
    ((0.4084, 0.3708, 0.1488, 1.0), 0.1809),
    ((0.3661, 0.3106, 0.1336, 1.0), 0.2128),
    ((0.6299, 0.5461, 0.2758, 1.0), 0.266),
    ((0.4535, 0.4015, 0.1772, 1.0), 0.2872),
    ((0.5818, 0.5749, 0.2734, 1.0), 0.3085),
    ((0.3063, 0.2933, 0.1009, 1.0), 0.3298),
    ((0.4449, 0.4316, 0.1853, 1.0), 0.3723),
    ((0.2814, 0.3075, 0.1075, 1.0), 0.4043),
    ((0.3078, 0.3569, 0.1202, 1.0), 0.4468),
    ((0.3034, 0.3073, 0.1076, 1.0), 0.5),
    ((0.1296, 0.183, 0.0342, 1.0), 0.5319),
    ((0.2191, 0.2778, 0.0793, 1.0), 0.5638),
    ((0.1282, 0.1593, 0.0378, 1.0), 0.5957),
    ((0.3567, 0.4744, 0.1591, 1.0), 0.6064),
    ((0.1719, 0.2404, 0.0449, 1.0), 0.6277),
    ((0.2168, 0.2639, 0.0755, 1.0), 0.7234),
    ((0.3594, 0.3372, 0.1472, 1.0), 0.7553),
    ((0.1611, 0.1665, 0.0481, 1.0), 0.7872),
    ((0.3238, 0.3161, 0.1372, 1.0), 0.8085),
    ((0.0813, 0.0699, 0.0184, 1.0), 0.8191),
    ((0.2095, 0.2009, 0.0635, 1.0), 0.8723),
    ((0.7108, 0.6485, 0.3301, 1.0), 0.8936),
    ((0.3915, 0.3563, 0.1493, 1.0), 0.9043),
    ((0.1694, 0.1873, 0.0553, 1.0), 0.9574),
    ((0.4335, 0.4064, 0.1848, 1.0), 0.9894),
    ((0.2195, 0.1946, 0.0561, 1.0), 1.0)]

pallete3 = [   ((0.3821, 0.4798, 0.0818, 1.0), 0.0),
    ((0.0858, 0.1301, 0.0193, 1.0), 0.0316),
    ((0.1099, 0.1785, 0.0071, 1.0), 0.0842),
    ((0.8128, 0.9357, 0.197, 1.0), 0.1053),
    ((0.1502, 0.2148, 0.0081, 1.0), 0.1158),
    ((0.391, 0.5164, 0.088, 1.0), 0.1789),
    ((0.0937, 0.136, 0.0077, 1.0), 0.2421),
    ((0.4853, 0.5725, 0.1079, 1.0), 0.2526),
    ((0.149, 0.2143, 0.0111, 1.0), 0.2947),
    ((0.7084, 0.8469, 0.3095, 1.0), 0.3053),
    ((0.0018, 0.008, 0.0, 1.0), 0.3158),
    ((0.9181, 0.9247, 0.4006, 1.0), 0.3263),
    ((0.1496, 0.1989, 0.0209, 1.0), 0.3579),
    ((0.18, 0.2798, 0.0095, 1.0), 0.4211),
    ((0.0971, 0.1541, 0.0068, 1.0), 0.4632),
    ((0.3364, 0.4352, 0.0719, 1.0), 0.4947),
    ((0.1105, 0.1727, 0.0078, 1.0), 0.5579),
    ((0.3493, 0.5321, 0.0443, 1.0), 0.6316),
    ((0.1035, 0.1735, 0.0031, 1.0), 0.6526),
    ((0.6578, 0.7528, 0.1652, 1.0), 0.6737),
    ((0.0797, 0.1192, 0.0077, 1.0), 0.6947),
    ((0.126, 0.191, 0.0137, 1.0), 0.7579),
    ((0.4031, 0.5242, 0.0569, 1.0), 0.8),
    ((0.9708, 0.9771, 0.524, 1.0), 0.8211),
    ((0.0765, 0.1203, 0.0085, 1.0), 0.8316),
    ((0.4879, 0.6171, 0.1039, 1.0), 0.8947),
    ((0.154, 0.2268, 0.0237, 1.0), 0.9158),
    ((0.8421, 0.9664, 0.3916, 1.0), 0.9263),
    ((0.192, 0.2556, 0.0265, 1.0), 0.9474),
    ((0.4264, 0.5593, 0.0509, 1.0), 0.9579),
    ((0.1199, 0.1892, 0.004, 1.0), 1.0)]

pallete4 = [   ((0.0273, 0.0802, 0.0382, 1.0), 0.0),
    ((0.0232, 0.0742, 0.0356, 1.0), 0.0323),
    ((0.1095, 0.2159, 0.1221, 1.0), 0.0645),
    ((0.0296, 0.0865, 0.0319, 1.0), 0.0968),
    ((0.0545, 0.1274, 0.0666, 1.0), 0.129),
    ((0.1106, 0.2139, 0.1332, 1.0), 0.1613),
    ((0.0342, 0.0948, 0.0369, 1.0), 0.2258),
    ((0.0513, 0.1195, 0.0612, 1.0), 0.2581),
    ((0.0182, 0.0558, 0.0222, 1.0), 0.2903),
    ((0.0242, 0.0654, 0.0343, 1.0), 0.3548),
    ((0.0545, 0.1144, 0.0742, 1.0), 0.3871),
    ((0.0195, 0.0503, 0.0235, 1.0), 0.4194),
    ((0.0186, 0.0652, 0.0218, 1.0), 0.4839),
    ((0.0453, 0.1304, 0.0593, 1.0), 0.5484),
    ((0.0844, 0.1845, 0.1329, 1.0), 0.5806),
    ((0.0395, 0.1022, 0.0545, 1.0), 0.6129),
    ((0.0683, 0.1462, 0.0793, 1.0), 0.6452),
    ((0.0265, 0.0726, 0.0265, 1.0), 0.7097),
    ((0.0638, 0.1345, 0.0624, 1.0), 0.7419),
    ((0.017, 0.0631, 0.0173, 1.0), 0.8065),
    ((0.0295, 0.0821, 0.0318, 1.0), 0.8387),
    ((0.0184, 0.0643, 0.0211, 1.0), 0.9032),
    ((0.0298, 0.0866, 0.0336, 1.0), 0.9355),
    ((0.0193, 0.0679, 0.0271, 1.0), 1.0)]

pallete5 = [   ((0.1712, 0.2776, 0.0465, 1.0), 0.0),
    ((0.0596, 0.1252, 0.0085, 1.0), 0.0732),
    ((0.1746, 0.2918, 0.0561, 1.0), 0.0976),
    ((0.0666, 0.1413, 0.007, 1.0), 0.122),
    ((0.1068, 0.2119, 0.0179, 1.0), 0.1463),
    ((0.0924, 0.1889, 0.0124, 1.0), 0.2195),
    ((0.1363, 0.2272, 0.0336, 1.0), 0.2927),
    ((0.0964, 0.1607, 0.0198, 1.0), 0.3171),
    ((0.0938, 0.1642, 0.0263, 1.0), 0.3902),
    ((0.0612, 0.117, 0.0137, 1.0), 0.4146),
    ((0.157, 0.2397, 0.0658, 1.0), 0.439),
    ((0.0213, 0.0497, 0.0023, 1.0), 0.4878),
    ((0.0487, 0.0734, 0.0116, 1.0), 0.5366),
    ((0.0366, 0.0546, 0.006, 1.0), 0.5854),
    ((0.2384, 0.2789, 0.107, 1.0), 0.6098),
    ((0.1384, 0.1946, 0.0482, 1.0), 0.6341),
    ((0.2384, 0.2789, 0.107, 1.0), 0.6585),
    ((0.0408, 0.0628, 0.0031, 1.0), 0.6829),
    ((0.0842, 0.1291, 0.0168, 1.0), 0.7317),
    ((0.0727, 0.1155, 0.0125, 1.0), 0.7561),
    ((0.1238, 0.1865, 0.0352, 1.0), 0.8049),
    ((0.0761, 0.1247, 0.016, 1.0), 0.8537),
    ((0.1389, 0.1985, 0.0443, 1.0), 0.878),
    ((0.0646, 0.1021, 0.0119, 1.0), 0.9512),
    ((0.1329, 0.1812, 0.0395, 1.0), 0.9756),
    ((0.1195, 0.1651, 0.0319, 1.0), 1.0)]

def shader_grass_texture_original(nw: NodeWrangler):
    # Code generated using version 2.4.3 of the node_transpiler

    texture_coordinate = nw.new_node(Nodes.TextureCoord)
    
    coord = nw.new_node(Nodes.VectorMath,
        input_kwargs={0: texture_coordinate.outputs["UV"], 'Scale': uniform(0.02, 0.2)},
        attrs={'operation': 'SCALE'})

    separate_xyz = nw.new_node(Nodes.SeparateXYZ,
        input_kwargs={'Vector': coord})
    
    edge_height = nw.new_node(Nodes.Math,
        input_kwargs={0: separate_xyz.outputs["X"], 1: 6.0, 2: -10.0},
        label='edge height',
        attrs={'operation': 'MULTIPLY_ADD'})
    
    combine_xyz = nw.new_node(Nodes.CombineXYZ,
        input_kwargs={'X': separate_xyz.outputs["Y"], 'Y': edge_height})
    
    wave_texture = nw.new_node(Nodes.WaveTexture,
        input_kwargs={'Vector': separate_xyz.outputs["Y"], 'Scale': 25.0, 'Distortion': 8.0, 'Detail Scale': 6.0})
    
    musgrave_texture = nw.new_node(Nodes.MusgraveTexture,
        input_kwargs={'Vector': combine_xyz, 'Scale': 8.0, 'Detail': 5.0, 'Dimension': 0.1, 'Lacunarity': 3.0})
    
    mix_1 = nw.new_node(Nodes.MixRGB,
        input_kwargs={'Fac': 0.2, 'Color1': wave_texture.outputs["Color"], 'Color2': musgrave_texture})
    
    
    object_info = nw.new_node(Nodes.ObjectInfo_Shader)
    map_range_1 = nw.new_node(Nodes.MapRange, input_kwargs={0: uniform(), 3: object_info.outputs["Random"], 4: mix_1})
    colorramp = nw.new_node(Nodes.ColorRamp, input_kwargs={'Fac': map_range_1.outputs["Result"]})
    
    pallete = np.random.choice([pallete1, pallete2, pallete3, pallete4, pallete5])
    np.random.shuffle(pallete)
    pallete = pallete[:np.random.randint(4, len(pallete))]
    for _ in range(len(pallete)-2):
        colorramp.color_ramp.elements.new(0)
    assert len(pallete) == len(colorramp.color_ramp.elements)
    for el, (rgba, pos) in zip(colorramp.color_ramp.elements, pallete):
        el.position = pos
        el.color = tuple(rgba)

    rough1, rough2 = uniform(0.2, 0.6, 2)
    roughness = nw.new_node(Nodes.MapRange, [mix_1, 0, 1, rough1, rough2])
    
    principled_bsdf = nw.new_node(Nodes.PrincipledBSDF,
        input_kwargs={'Base Color': colorramp.outputs["Color"], 'Roughness': roughness})
    translucent = nw.new_node(Nodes.TranslucentBSDF, [colorramp.outputs["Color"]])
    shader = nw.new_node(Nodes.MixShader, [0.7, principled_bsdf, translucent])
    
    nw.new_node(Nodes.MaterialOutput, input_kwargs={'Surface': shader})



def apply(obj, selection=None, **kwargs):
    surface.add_material(obj, shader_grass_texture_original, selection=selection)